import os
import logging
import time
import glob
import json
import requests
import uuid
from datetime import datetime
from typing import List, Dict, Optional
from dataclasses import dataclass
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.options import Options
from sqlalchemy import update
from app.db.models import TPost
from app.db.database import SessionLocal
from selenium.common.exceptions import TimeoutException, NoSuchElementException

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('tts_automation.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

@dataclass
class Article:
    id: int
    content: str

class TTSAutomation:
    def __init__(self, base_dir: str = "/app/downloads", username: str = "zxsr20", password: str = "aA@#$67028"):
        self.base_dir = base_dir
        self.username = username
        self.password = password
        self.driver = None
        self.is_logged_in = False
        self.cookie_file = os.path.join(base_dir, "tts_cookies.json")
        self.selenium_hub_url = 'http://selenium-hub:4444'
        
        # 创建基础下载目录
        os.makedirs(self.base_dir, exist_ok=True)
        
        # 创建今天的日期目录
        self.today_dir = os.path.join(self.base_dir, datetime.now().strftime('%Y%m%d'))
        os.makedirs(self.today_dir, exist_ok=True)

    def _cleanup_selenium_sessions(self):
        """清理所有活跃的Selenium会话"""
        try:
            # 获取所有活跃的会话
            response = requests.get(f"{self.selenium_hub_url}/wd/hub/sessions")
            if response.status_code == 200:
                sessions = response.json().get('value', [])
                for session in sessions:
                    session_id = session.get('id')
                    if session_id:
                        # 删除每个会话
                        try:
                            delete_url = f"{self.selenium_hub_url}/wd/hub/session/{session_id}"
                            requests.delete(delete_url)
                            logger.info(f"Cleaned up session: {session_id}")
                        except Exception as e:
                            logger.error(f"Error cleaning up session {session_id}: {str(e)}")
        except Exception as e:
            logger.error(f"Error getting Selenium sessions: {str(e)}")

    def _setup_driver(self):
        """设置并初始化WebDriver"""
        try:
            # 首先清理现有会话
            self._cleanup_selenium_sessions()
            
            # 如果已经有driver实例，先关闭它
            if self.driver:
                try:
                    self.driver.quit()
                except:
                    pass
                self.driver = None
            
            chrome_options = Options()
            chrome_options.add_argument('--no-sandbox')
            chrome_options.add_argument('--disable-dev-shm-usage')
            chrome_options.add_argument('--headless')
            chrome_options.add_argument('--disable-gpu')
            
            prefs = {
                "download.default_directory": "/home/seluser/Downloads",
                "download.prompt_for_download": False,
                "download.directory_upgrade": True,
                "safebrowsing.enabled": True,
                "download.default_directory_infobar_shown": False
            }
            chrome_options.add_experimental_option("prefs", prefs)
            chrome_options.set_capability("platformName", "linux")
            
            logger.info("Connecting to Selenium Grid")
            self.driver = webdriver.Remote(
                command_executor=f'{self.selenium_hub_url}/wd/hub',
                options=chrome_options
            )
            logger.info("Successfully connected to Selenium Grid")
        except Exception as e:
            logger.error(f"Error setting up WebDriver: {str(e)}")
            raise

    def _save_cookies(self):
        """保存cookies到文件"""
        if self.driver:
            cookies = self.driver.get_cookies()
            with open(self.cookie_file, 'w') as f:
                json.dump(cookies, f)
            logger.info("Cookies saved successfully")

    def _load_cookies(self) -> bool:
        """从文件加载cookies"""
        try:
            if os.path.exists(self.cookie_file):
                with open(self.cookie_file, 'r') as f:
                    cookies = json.load(f)
                for cookie in cookies:
                    try:
                        self.driver.add_cookie(cookie)
                    except Exception as e:
                        logger.error(f"Error adding cookie: {str(e)}")
                logger.info("Cookies loaded successfully")
                return True
            return False
        except Exception as e:
            logger.error(f"Error loading cookies: {str(e)}")
            return False

    def _login(self) -> bool:
        """执行登录操作"""
        try:
            logger.info("Starting login process")
            self.driver.get("http://new.text-to-speech.cn/tts/")
            
            # 尝试加载cookies
            if self._load_cookies():
                # 刷新页面使cookies生效
                self.driver.refresh()
                # 检查是否已登录
                self._handle_popup()
                try:

                    # 等待按钮可见和可点击
                    WebDriverWait(self.driver, 10).until(
                        EC.element_to_be_clickable((By.CSS_SELECTOR, "button#btnUpdate"))
                    )
                    logger.info("Already logged in via cookies")
                    self.is_logged_in = True
                    return True
                except:
                    logger.info("Cookies expired or invalid, proceeding with manual login")
            
            # 如果cookie登录失败，执行手动登录
            nav_login_button = WebDriverWait(self.driver, 10).until(
                EC.element_to_be_clickable((By.CSS_SELECTOR, "a.login-btn.navbar-button"))
            )
            nav_login_button.click()
            
            # 输入用户名和密码
            username_input = WebDriverWait(self.driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, "input[placeholder='请输入电子邮箱/用户名']"))
            )
            username_input.send_keys(self.username)
            
            password_input = self.driver.find_element(By.CSS_SELECTOR, "input[placeholder='请输入密码']")
            password_input.send_keys(self.password)
            
            # 点击登录按钮
            login_button = self.driver.find_element(By.CSS_SELECTOR, "button.go-login")
            self.driver.execute_script("arguments[0].click();", login_button)
            
            # 等待登录成功
            time.sleep(3)
            self.is_logged_in = True
            
            # 保存cookies以供下次使用
            self._save_cookies()
            
            logger.info("Login successful")
            return True
            
        except Exception as e:
            logger.error(f"Login failed: {str(e)}")
            return False

    def _wait_for_download(self, timeout: int = 30) -> Optional[str]:
        """等待文件下载完成并返回文件路径"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            # 检查下载目录中的所有文件，只搜索文件而不是目录
            try:
                files = [f for f in glob.glob(os.path.join(self.base_dir, '*')) 
                        if os.path.isfile(f) and not f.endswith('.crdownload') 
                        and not f.endswith('.tmp') 
                        and not f.endswith('tts_cookies.json')]  # 排除cookie文件
                
                if files:
                    # 按修改时间排序，获取最新的文件
                    latest_file = max(files, key=os.path.getmtime)
                    logger.info(f"Found file: {latest_file}")
                    
                    # 检查文件大小和类型
                    if os.path.getsize(latest_file) > 0 and latest_file.endswith('.mp3'):
                        return latest_file
                    else:
                        logger.error(f"Invalid file: {latest_file} (size: {os.path.getsize(latest_file)})")
                else:
                    logger.info("No valid files found in download directory yet...")
            except Exception as e:
                logger.error(f"Error while checking download directory: {str(e)}")
            
            time.sleep(1)
            
        logger.error(f"Download timeout after {timeout} seconds")
        return None

    def _handle_popup(self):
        """处理页面上的弹窗"""
        try:
            # 等待弹窗出现
            popup = WebDriverWait(self.driver, 5).until(
                EC.presence_of_element_located((By.CLASS_NAME, "swal2-popup"))
            )
            
            if popup.is_displayed():
                # 首先尝试查找并点击关闭按钮（×）
                try:
                    close_button = WebDriverWait(self.driver, 3).until(
                        EC.element_to_be_clickable((By.CLASS_NAME, "swal2-close"))
                    )
                    logger.info("找到关闭按钮，点击关闭")
                    self.driver.execute_script("arguments[0].click();", close_button)
                    time.sleep(1)  # 等待弹窗消失
                    return True
                except:
                    logger.info("未找到关闭按钮，尝试确认按钮")
                
                # 如果没有关闭按钮，尝试确认按钮
                try:
                    confirm_button = WebDriverWait(self.driver, 3).until(
                        EC.element_to_be_clickable((By.CLASS_NAME, "swal2-confirm"))
                    )
                    logger.info("找到确认按钮，点击确认")
                    self.driver.execute_script("arguments[0].click();", confirm_button)
                    time.sleep(1)  # 等待弹窗消失
                    return True
                except:
                    logger.warning("未找到确认按钮")
                
                # 如果都没有找到，尝试直接通过JavaScript关闭弹窗
                try:
                    logger.info("尝试通过JavaScript关闭弹窗")
                    self.driver.execute_script("""
                        var elements = document.getElementsByClassName('swal2-container');
                        for(var i=0; i<elements.length; i++) {
                            elements[i].remove();
                        }
                    """)
                    time.sleep(1)
                    return True
                except:
                    logger.warning("JavaScript关闭弹窗失败")
                
                return False
                
        except TimeoutException:
            # 没有找到弹窗，说明可能没有弹窗
            return False
        except Exception as e:
            logger.warning(f"处理弹窗时出现异常: {str(e)}")
            return False
        return False

    def _process_single_article(self, article: Article) -> Optional[str]:
        """处理单篇文章的转换"""
        try:
            # 检查文章内容是否为空
            if not article.content or len(article.content.strip()) == 0:
                logger.info(f"Skipping article {article.id} due to empty content")
                return None
                
            # 等待文本框加载并输入文本
            text_area = WebDriverWait(self.driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, "textarea.form-control"))
            )
            text_area.clear()
            text_area.send_keys(article.content)
            logger.info(f"article content: {article.content}")
            
            max_attempts = 3
            for attempt in range(max_attempts):
                try:
                    # 处理可能存在的弹窗
                    self._handle_popup()
                    
                    # 等待按钮可见和可点击
                    generate_button = WebDriverWait(self.driver, 10).until(
                        EC.element_to_be_clickable((By.CSS_SELECTOR, "button#btnUpdate"))
                    )
                    
                    # 确保按钮在视图中
                    self.driver.execute_script("arguments[0].scrollIntoView(true);", generate_button)
                    time.sleep(1)  # 等待滚动完成
                    
                    logger.info("点击生成按钮")
                    self.driver.execute_script("arguments[0].click();", generate_button)
                    logger.info("已点击生成按钮")
                    break
                except TimeoutException:
                    # 打印当前页面源码和状态
                    logger.warning(f"第 {attempt + 1} 次尝试点击生成按钮失败，当前页面状态：")
                    try:
                        # 尝试直接找到按钮（不等待）看看是否存在
                        button = self.driver.find_element(By.CSS_SELECTOR, "button#btnUpdate")
                        logger.warning(f"按钮存在但不可点击，按钮HTML: {button.get_attribute('outerHTML')}")
                        logger.warning(f"按钮是否显示: {button.is_displayed()}")
                        logger.warning(f"按钮是否启用: {button.is_enabled()}")
                    except:
                        logger.warning("按钮完全不存在于页面中")
                    
                    # 打印整个页面源码
                    # logger.warning(f"页面源码: {self.driver.page_source}")
                    
                    if attempt == max_attempts - 1:
                        raise Exception("无法点击生成按钮")
                    logger.warning(f"重试中...")
                    time.sleep(2)
            
            # 等待生成完成 - 检查下载按钮的状态
            try:
                # 等待下载按钮出现并可点击
                download_button = None
                download_selectors = [
                    "//button[contains(text(), '下载')]",
                    "//button[contains(@class, 'download')]",
                    "//button[contains(@class, 'btn-success')]",
                    "//a[contains(text(), '下载')]"
                ]
                
                for selector in download_selectors:
                    try:
                        download_button = WebDriverWait(self.driver, 60).until(
                            EC.element_to_be_clickable((By.XPATH, selector))
                        )
                        if download_button:
                            logger.info(f"Found clickable download button with selector: {selector}")
                            break
                    except:
                        continue
                
                if not download_button:
                    raise Exception("Download button not found or not clickable")
                
                # 点击下载按钮
                self.driver.execute_script("arguments[0].click();", download_button)
                
            except Exception as e:
                logger.error(f"Error waiting for generation to complete: {str(e)}")
                # 记录当前页面状态
                try:
                    logger.error(f"Current page source: {self.driver.page_source}")
                except:
                    logger.error("Could not get page source")
                raise Exception("TTS generation timeout or failed")
            
            # 等待下载完成
            downloaded_file = self._wait_for_download()
            if not downloaded_file:
                raise Exception("Download failed or timed out")
            
            logger.info(f"Downloaded file path: {downloaded_file}")
            
            # 使用UUID生成唯一文件名
            file_uuid = str(uuid.uuid4())
            new_filename = f"{file_uuid}.mp3"
            new_filepath = os.path.join(self.today_dir, new_filename)
            
            # 如果目标文件已存在，先删除
            if os.path.exists(new_filepath):
                logger.info(f"Removing existing file: {new_filepath}")
                os.remove(new_filepath)
            
            logger.info(f"Renaming {downloaded_file} to {new_filepath}")
            os.rename(downloaded_file, new_filepath)
            
            # 验证文件是否成功重命名
            if os.path.exists(new_filepath):
                logger.info(f"Successfully processed article {article.id}")
                with SessionLocal() as db:
                    stmt = update(TPost).where(TPost.id == article.id).values(ttl=new_filepath)
                    db.execute(stmt)
                    db.commit()
                logger.info(f"Article {article.id} -> {new_filepath}")
                return new_filepath
            else:
                raise Exception(f"Failed to rename file to {new_filepath}")
            
        except Exception as e:
            logger.error(f"Error processing article {article.id}: {str(e)}")
            logger.error(f"Current directory contents: {os.listdir(self.base_dir)}")
            return None

    def process_articles(self, articles: List[Article]) -> Dict[int, str]:
        """批量处理文章列表"""
        results = {}
        
        try:
            # 过滤掉空内容的文章
            valid_articles = [article for article in articles if article.content and len(article.content.strip()) > 0]
            
            if not valid_articles:
                logger.info("No valid articles to process (all articles have empty content)")
                return results
                
            logger.info(f"Processing {len(valid_articles)} articles out of {len(articles)} total")
            
            self._setup_driver()
            if not self._login():
                raise Exception("Failed to login")
            
            for article in valid_articles:
                filepath = self._process_single_article(article)
                if filepath:
                    results[article.id] = filepath
                
        except Exception as e:
            logger.error(f"Error in batch processing: {str(e)}")
        finally:
            if self.driver:
                self.driver.quit()
        
        return results

    def test_tts(self):
        # 测试代码
        test_articles = [
            Article(id=1, content="这是第一篇测试文章的内容。"),
            Article(id=2, content="这是第二篇测试文章的内容。")
        ]
        results = self.process_articles(test_articles)
        for article_id, filepath in results.items():
            print(f"Article {article_id} -> {filepath}") 
        return results

    def __del__(self):
        """析构函数，确保在对象被销毁时关闭driver"""
        if self.driver:
            try:
                self.driver.quit()
            except:
                pass

if __name__ == "__main__":
    # 测试代码
    test_articles = [
        Article(id=1, content="这是第一篇测试文章的内容。"),
        Article(id=2, content="这是第二篇测试文章的内容。")
    ]
    
    tts = TTSAutomation()
    results = tts.process_articles(test_articles)
    
    for article_id, filepath in results.items():
        print(f"Article {article_id} -> {filepath}") 
    